import tensorflow as tf
from tensorflow.keras import layers, models


class CNNModel:
    def __init__(self):
        # 在初始化时构建模型
        self.model = models.Sequential([
            layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
            layers.MaxPooling2D((2, 2)),
            layers.Conv2D(64, (3, 3), activation='relu'),
            layers.MaxPooling2D((2, 2)),
            layers.Conv2D(64, (3, 3), activation='relu'),
            layers.Flatten(),
            layers.Dense(64, activation='relu'),
            layers.Dense(10, activation='softmax')
        ])

    def summary(self, ):
        # 模型概要打印
        self.model.summary()

    def compile(self, optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy']):
        # 编译模型
        self.model.compile(optimizer=optimizer, loss=loss, metrics=metrics)

    def train(self, train_images, train_labels, epochs=5, batch_size=64, validation_data=None):
        # 训练模型
        return self.model.fit(train_images, train_labels, epochs=epochs, batch_size=batch_size,
                              validation_data=validation_data)

    def evaluate(self, test_images, test_labels, verbose=2):
        # 评估模型
        return self.model.evaluate(test_images, test_labels, verbose=verbose)

    def save(self, path):
        # 保存模型
        self.model.save(path)

    @staticmethod
    def load(path):
        # 静态方法：加载模型
        return tf.keras.models.load_model(path)
